{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspecting TFX metadata\n",
    "\n",
    "\n",
    "## Learning Objectives\n",
    "\n",
    "1. Use a GRPC server to access and analyze pipeline artifacts stored in the ML Metadata service of your AI Platform Pipelines instance.\n",
    "\n",
    "In this lab, you will explore TFX pipeline metadata including pipeline and run artifacts. A hosted **AI Platform Pipelines** instance includes the [ML Metadata](https://github.com/google/ml-metadata) service. In **AI Platform Pipelines**, ML Metadata uses *MySQL* as a database backend and can be accessed using a GRPC server."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "import ml_metadata\n",
    "import tensorflow_data_validation as tfdv\n",
    "import tensorflow_model_analysis as tfma\n",
    "\n",
    "\n",
    "from ml_metadata.metadata_store import metadata_store\n",
    "from ml_metadata.proto import metadata_store_pb2\n",
    "\n",
    "from tfx.orchestration import metadata\n",
    "from tfx.types import standard_artifacts\n",
    "\n",
    "from tensorflow.python.lib.io import file_io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import tfx; print('TFX version: {}'.format(tfx.__version__))\"\n",
    "!python -c \"import kfp; print('KFP version: {}'.format(kfp.__version__))\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 1: Explore metadata from existing TFX pipeline runs from AI Pipelines instance created in `lab-02` or `lab-03`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.1 Configure Kubernetes port forwarding\n",
    "\n",
    "To enable access to the ML Metadata GRPC server, configure Kubernetes port forwarding.\n",
    "\n",
    "From a JupyterLab terminal, execute the following commands:\n",
    "\n",
    "```\n",
    "gcloud container clusters get-credentials [YOUR CLUSTER] --zone [YOUR CLUSTER ZONE]  \n",
    "kubectl port-forward  service/metadata-grpc-service --namespace [YOUR NAMESPACE] 7000:8080\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Proceed to the next step, \"Connecting to ML Metadata\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2: Create new AI Pipelines instance and evaluate metadata on newly triggered pipeline runs.\n",
    "\n",
    "Hosted AI Pipelines incurs cost for the duration your Kubernetes cluster is running. If you deleted your previous lab instance, proceed with the 6 steps below to deploy a new TFX pipeline and triggers runs to inspect its metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "\n",
    "# Set `PATH` to include the directory containing TFX CLI.\n",
    "PATH=%env PATH\n",
    "%env PATH=/home/jupyter/.local/bin:{PATH}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pipeline source can be found in the `pipeline` folder. Switch to the `pipeline` folder and compile the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1 Create AI Platform Pipelines cluster\n",
    "\n",
    "Navigate to [AI Platform Pipelines](https://console.cloud.google.com/ai-platform/pipelines/clusters) page in the Google Cloud Console.\n",
    "\n",
    "Create or select an existing Kubernetes cluster (GKE) and deploy AI Platform. Make sure to select `\"Allow access to the following Cloud APIs https://www.googleapis.com/auth/cloud-platform\"` to allow for programmatic access to your pipeline by the Kubeflow SDK for the rest of the lab. Also, provide an `App instance name` such as \"TFX-lab-04\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2 Configure environment settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update  the below constants  with the settings reflecting your lab environment.\n",
    "\n",
    "- `GCP_REGION` - the compute region for AI Platform Training and Prediction\n",
    "- `ARTIFACT_STORE` - the GCS bucket created during installation of AI Platform Pipelines. The bucket name starts with the `kubeflowpipelines-` prefix. Alternatively, you can specify create a new storage bucket to write pipeline artifacts to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gsutil ls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `CUSTOM_SERVICE_ACCOUNT` - In the gcp console Click on the Navigation Menu. Navigate to `IAM & Admin`, then to `Service Accounts` and use the service account starting with prifix - `'tfx-tuner-caip-service-account'`. This enables CloudTuner and the Google Cloud AI Platform extensions Tuner component to work together and allows for distributed and parallel tuning backed by AI Platform Vizier's hyperparameter search algorithm. Please see the lab setup `README` for setup instructions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `ENDPOINT` - set the `ENDPOINT` constant to the endpoint to your AI Platform Pipelines instance. The endpoint to the AI Platform Pipelines instance can be found on the [AI Platform Pipelines](https://console.cloud.google.com/ai-platform/pipelines/clusters) page in the Google Cloud Console.\n",
    "\n",
    "1. Open the *SETTINGS* for your instance\n",
    "2. Use the value of the `host` variable in the *Connect to this Kubeflow Pipelines instance from a Python client via Kubeflow Pipelines SKD* section of the *SETTINGS* window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: Set your environment resource settings here for GCP_REGION, ARTIFACT_STORE_URI, ENDPOINT, and CUSTOM_SERVICE_ACCOUNT.\n",
    "GCP_REGION = 'us-central1'\n",
    "ARTIFACT_STORE_URI = 'gs://dougkelly-sandbox-kubeflowpipelines-default' #Change\n",
    "ENDPOINT = '60ff837483ecde05-dot-us-central2.pipelines.googleusercontent.com' #Change\n",
    "CUSTOM_SERVICE_ACCOUNT = 'tfx-tuner-caip-service-account@dougkelly-sandbox.iam.gserviceaccount.com' #Change\n",
    "\n",
    "PROJECT_ID = !(gcloud config get-value core/project)\n",
    "PROJECT_ID = PROJECT_ID[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set your resource settings as environment variables. These override the default values in pipeline/config.py.\n",
    "%env GCP_REGION={GCP_REGION}\n",
    "%env ARTIFACT_STORE_URI={ARTIFACT_STORE_URI}\n",
    "%env CUSTOM_SERVICE_ACCOUNT={CUSTOM_SERVICE_ACCOUNT}\n",
    "%env PROJECT_ID={PROJECT_ID}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3 Compile pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PIPELINE_NAME = 'tfx_covertype_lab_04'\n",
    "MODEL_NAME = 'tfx_covertype_classifier'\n",
    "DATA_ROOT_URI = 'gs://workshop-datasets/covertype/small'\n",
    "CUSTOM_TFX_IMAGE = 'gcr.io/{}/{}'.format(PROJECT_ID, PIPELINE_NAME)\n",
    "RUNTIME_VERSION = '2.3'\n",
    "PYTHON_VERSION = '3.7'\n",
    "USE_KFP_SA=False\n",
    "ENABLE_TUNING=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env PIPELINE_NAME={PIPELINE_NAME}\n",
    "%env MODEL_NAME={MODEL_NAME}\n",
    "%env DATA_ROOT_URI={DATA_ROOT_URI}\n",
    "%env KUBEFLOW_TFX_IMAGE={CUSTOM_TFX_IMAGE}\n",
    "%env RUNTIME_VERSION={RUNTIME_VERSION}\n",
    "%env PYTHON_VERIONS={PYTHON_VERSION}\n",
    "%env USE_KFP_SA={USE_KFP_SA}\n",
    "%env ENABLE_TUNING={ENABLE_TUNING}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tfx pipeline compile --engine kubeflow --pipeline_path runner.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.4 Deploy pipeline to AI Platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tfx pipeline create  \\\n",
    "--pipeline_path=runner.py \\\n",
    "--endpoint={ENDPOINT} \\\n",
    "--build_target_image={CUSTOM_TFX_IMAGE}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(optional) If you make local changes to the pipeline, you can update the deployed package on AI Platform with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tfx pipeline update --pipeline_path runner.py --endpoint {ENDPOINT}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.5 Create and monitor pipeline run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tfx run create --pipeline_name={PIPELINE_NAME} --endpoint={ENDPOINT}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.6 Configure Kubernetes port forwarding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To enable access to the ML Metadata GRPC server, configure Kubernetes port forwarding.\n",
    "\n",
    "From a JupyterLab terminal, execute the following commands:\n",
    "\n",
    "```\n",
    "gcloud container clusters get-credentials [YOUR CLUSTER] --zone [YOURE CLUSTER ZONE]  \n",
    "kubectl port-forward  service/metadata-grpc-service --namespace [YOUR NAMESPACE] 7000:8080\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to ML Metadata "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure ML Metadata GRPC client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grpc_host = 'localhost'\n",
    "grpc_port = 7000\n",
    "connection_config = metadata_store_pb2.MetadataStoreClientConfig()\n",
    "connection_config.host = grpc_host\n",
    "connection_config.port = grpc_port"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connect to ML Metadata service"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "store = metadata_store.MetadataStore(connection_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Important\n",
    "\n",
    "A full pipeline run without tuning takes about 40-45 minutes to complete. You need to wait until a pipeline run is complete before proceeding with the steps below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploring ML Metadata \n",
    "\n",
    "The Metadata Store uses the following data model:\n",
    "\n",
    "- `ArtifactType` describes an artifact's type and its properties that are stored in the Metadata Store. These types can be registered on-the-fly with the Metadata Store in code, or they can be loaded in the store from a serialized format. Once a type is registered, its definition is available throughout the lifetime of the store.\n",
    "- `Artifact` describes a specific instances of an ArtifactType, and its properties that are written to the Metadata Store.\n",
    "- `ExecutionType` describes a type of component or step in a workflow, and its runtime parameters.\n",
    "- `Execution` is a record of a component run or a step in an ML workflow and the runtime parameters. An Execution can be thought of as an instance of an ExecutionType. Every time a developer runs an ML pipeline or step, executions are recorded for each step.\n",
    "- `Event` is a record of the relationship between an Artifact and Executions. When an Execution happens, Events record every Artifact that was used by the Execution, and every Artifact that was produced. These records allow for provenance tracking throughout a workflow. By looking at all Events MLMD knows what Executions happened, what Artifacts were created as a result, and can recurse back from any Artifact to all of its upstream inputs.\n",
    "- `ContextType` describes a type of conceptual group of Artifacts and Executions in a workflow, and its structural properties. For example: projects, pipeline runs, experiments, owners.\n",
    "- `Context` is an instances of a ContextType. It captures the shared information within the group. For example: project name, changelist commit id, experiment annotations. It has a user-defined unique name within its ContextType.\n",
    "- `Attribution` is a record of the relationship between Artifacts and Contexts.\n",
    "- `Association` is a record of the relationship between Executions and Contexts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "List the registered artifact types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for artifact_type in store.get_artifact_types():\n",
    "    print(artifact_type.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Display the registered execution types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for execution_type in store.get_execution_types():\n",
    "    print(execution_type.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "List the registered context types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for context_type in store.get_context_types():\n",
    "    print(context_type.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing TFX artifacts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve data analysis and validation artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with metadata.Metadata(connection_config) as store:\n",
    "    schema_artifacts = store.get_artifacts_by_type(standard_artifacts.Schema.TYPE_NAME)    \n",
    "    stats_artifacts = store.get_artifacts_by_type(standard_artifacts.ExampleStatistics.TYPE_NAME)\n",
    "    anomalies_artifacts = store.get_artifacts_by_type(standard_artifacts.ExampleAnomalies.TYPE_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_file = os.path.join(schema_artifacts[-1].uri, 'schema.pbtxt')\n",
    "print(\"Generated schame file:{}\".format(schema_file))\n",
    "\n",
    "stats_path = stats_artifacts[-1].uri\n",
    "train_stats_file = os.path.join(stats_path, 'train', 'stats_tfrecord')\n",
    "eval_stats_file = os.path.join(stats_path, 'eval', 'stats_tfrecord')\n",
    "print(\"Train stats file:{}, Eval stats file:{}\".format(\n",
    "    train_stats_file, eval_stats_file))\n",
    "\n",
    "anomalies_path = anomalies_artifacts[-1].uri\n",
    "train_anomalies_file = os.path.join(anomalies_path, 'train', 'anomalies.pbtxt')\n",
    "eval_anomalies_file = os.path.join(anomalies_path, 'eval', 'anomalies.pbtxt')\n",
    "\n",
    "print(\"Train anomalies file:{}, Eval anomalies file:{}\".format(\n",
    "    train_anomalies_file, eval_anomalies_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = tfdv.load_schema_text(schema_file)\n",
    "tfdv.display_schema(schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise: looking at the features visualized below, answer the following questions:\n",
    "\n",
    "- Which feature transformations would you apply to each feature with TF Transform?\n",
    "- Are there data quality issues with certain features that may impact your model performance? How might you deal with it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_stats = tfdv.load_statistics(train_stats_file)\n",
    "eval_stats = tfdv.load_statistics(eval_stats_file)\n",
    "tfdv.visualize_statistics(lhs_statistics=eval_stats, rhs_statistics=train_stats,\n",
    "                          lhs_name='EVAL_DATASET', rhs_name='TRAIN_DATASET')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_anomalies = tfdv.load_anomalies_text(train_anomalies_file)\n",
    "tfdv.display_anomalies(train_anomalies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_anomalies = tfdv.load_anomalies_text(eval_anomalies_file)\n",
    "tfdv.display_anomalies(eval_anomalies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve model artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with metadata.Metadata(connection_config) as store:\n",
    "    model_eval_artifacts = store.get_artifacts_by_type(standard_artifacts.ModelEvaluation.TYPE_NAME)\n",
    "    hyperparam_artifacts = store.get_artifacts_by_type(standard_artifacts.HyperParameters.TYPE_NAME)\n",
    "    \n",
    "model_eval_path = model_eval_artifacts[-1].uri\n",
    "print(\"Generated model evaluation result:{}\".format(model_eval_path))\n",
    "best_hparams_path = os.path.join(hyperparam_artifacts[-1].uri, 'best_hyperparameters.txt')\n",
    "print(\"Generated model best hyperparameters result:{}\".format(best_hparams_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Return best hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Latest pipeline run Tuner search space.\n",
    "json.loads(file_io.read_file_to_string(best_hparams_path))['space']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Latest pipeline run Tuner searched best_hyperparameters artifacts.\n",
    "json.loads(file_io.read_file_to_string(best_hparams_path))['values']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize model evaluations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise: review the model evaluation results below and answer the following questions:\n",
    "\n",
    "- Which Wilderness Area had the highest accuracy?\n",
    "- Which Wilderness Area had the lowest performance? Why do you think that is? What are some steps you could take to improve your next model runs?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_result = tfma.load_eval_result(model_eval_path)\n",
    "tfma.view.render_slicing_metrics(\n",
    "    eval_result, slicing_column='Wilderness_Area')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Debugging tip**: If the TFMA visualization of the Evaluator results do not render, try switching to view in a Classic Jupyter Notebook. You do so by clicking `Help > Launch Classic Notebook` and re-opening the notebook and running the above cell to see the interactive TFMA results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## License"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font size=-1>Licensed under the Apache License, Version 2.0 (the \\\"License\\\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \\\"AS IS\\\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the License for the specific language governing permissions and limitations under the License.</font>\n"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-gpu.2-3.m61",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-3:m61"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
